#include "SNCascadeClassifier.h"
//---------------------------------------------------

SNCascadeClassifier::SNCascadeClassifier()
{
}
//---------------------------------------------------

SNCascadeClassifier::~SNCascadeClassifier()
{
}

bool SNCascadeClassifier::Train(const QString cascade_dir_name, const QString pos_filename, const QString neg_filename, int num_pos, int num_neg, int precalc_val_bufsize, int precalc_idx_bufsize, int num_stages, const SNCascadeParams& cascade_params, const SNFeatureParams& feature_params, const SNCascadeBoostParams& stage_params, bool baseFormatSave /*= false*/)
{
	NumPos = num_pos;
	NumNeg = num_neg;
	NumStages = num_stages;
	if (!ImgReader.Create(pos_filename, neg_filename, cascade_params.WinSize))
	{
		//cout << "Image reader can not be created from -vec " << _posFilename
			//<< " and -bg " << _negFilename << "." << endl;
		return false;
	}

	CascadeParams = cascade_params;
	FeatureParams = SNFeatureParams::Create(cascade_params.FeatureType);

	FeatureParams->Init(feature_params);

	StageParams = new SNCascadeBoostParams;
	*StageParams = stage_params;
	FeatureEvaluator = SNFeatureEvaluator::Create(CascadeParams.FeatureType);
	FeatureEvaluator->Init((SNFeatureParams*)FeatureParams, NumPos + NumNeg, CascadeParams.WinSize);
	StageClassifiers.reserve(num_stages);

	int startNumStages = (int)StageClassifiers.size();
	if (startNumStages > 1)
		std::cout << std::endl << "Stages 0-" << startNumStages - 1 << " are loaded" << std::endl;
	else if (startNumStages == 1)
		std::cout << std::endl << "Stage 0 is loaded" << std::endl;

	double required_leaf_FA_rate = pow((double)StageParams->MaxFalseAlarm, (double)NumStages) / (double)StageParams->max_depth;

	double temp_leaf_FA_rate;

	for (int i = startNumStages; i < NumStages; i++)
	{
		std::cout << std::endl << "===== TRAINING " << i << "-stage =====" << std::endl;
		std::cout << "<BEGIN" << std::endl;

		if (!UpdateTrainingSet(temp_leaf_FA_rate))
		{
			std::cout << "Train dataset for temp stage can not be filled. "
				"Branch training terminated." << std::endl;
			break;
		}

		if (temp_leaf_FA_rate <= required_leaf_FA_rate)
		{
			std::cout << "Required leaf false alarm rate achieved. "
				"Branch training terminated." << std::endl;
			break;
		}

		SNCascadeBoost* temp_stage = new SNCascadeBoost;
		bool is_stage_trained = temp_stage->Train((SNFeatureEvaluator*)FeatureEvaluator, CurNumSamples, precalc_val_bufsize, precalc_idx_bufsize, *((SNCascadeBoostParams*)StageParams));
		std::cout << "END>" << std::endl;

		if (!is_stage_trained)
			break;
	}

	return true;
}
//---------------------------------------------------

bool SNCascadeClassifier::UpdateTrainingSet(double& acceptance_ratio)
{
	int64 pos_consumed = 0;
	int64 neg_consumed = 0;
	ImgReader.Restart();

	int pos_count = FillPassedSamples(0, NumPos, true, pos_consumed);

	if (!pos_count)
		return false;
	
	std::cout << "POS count : consumed   " << pos_count << " : " << (int)pos_consumed << std::endl;

	// apply only a fraction of negative samples. double is required since overflow is possible
	int pro_num_neg = cvRound((((double)NumNeg) * ((double)pos_count)) / NumPos); 

	int neg_count = FillPassedSamples(pos_count, pro_num_neg, false, neg_consumed);

	if (!neg_count)
		return false;

	CurNumSamples = pos_count + neg_count;

	acceptance_ratio = neg_consumed == 0 ? 0 : ((double)neg_count / (double)(int64)neg_consumed);

	std::cout << "NEG count : acceptanceRatio    " << neg_count << " : " << acceptance_ratio << std::endl;
	return true;
}
//---------------------------------------------------

int SNCascadeClassifier::FillPassedSamples(int first, int count, bool is_positive, int64& consumed)
{
	int get_count = 0;

	cv::Mat img(CascadeParams.WinSize, CV_8UC1);

	for (int i = first; i < first + count; i++)
	{
		for (;;)
		{
			bool is_get_img = is_positive ? ImgReader.GetPos(img) : ImgReader.GetNeg(img);

			if (!is_get_img)
				return get_count;

			consumed++;

			FeatureEvaluator->SetImage(img, is_positive ? 1 : 0, i);

			if (Predict(i) == 1.0F)
			{
				get_count++;
				printf("%s current samples: %d\r", is_positive ? "POS" : "NEG", get_count);
				break;
			}
		}
	}
	return get_count;
}
//---------------------------------------------------

int SNCascadeClassifier::Predict(int sample_idx)
{
	//CV_DbgAssert(sampleIdx < numPos + numNeg);
	for (std::vector< cv::Ptr<SNCascadeBoost> >::iterator it = StageClassifiers.begin(); it != StageClassifiers.end(); it++)
	{
		if ((*it)->Predict(sample_idx) == 0.f)
			return 0;
	}

	return 1;
}
//---------------------------------------------------
